package org.x.chat.config;


import cn.dev33.satoken.stp.StpUtil;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.stereotype.Component;
import org.x.chat.until.SessionKit;

@Component
public class WebSocketAuthInterceptor implements ChannelInterceptor {

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

        if (accessor != null && StompCommand.CONNECT.equals(accessor.getCommand())) {

            String sessionId = accessor.getSessionId();
            // 从header获取token
            String token = accessor.getFirstNativeHeader("Authorization");
            if (token != null && token.startsWith("Bearer ")) {
                token = token.substring(7);
            }
            // SaToken验证
            if (token == null) {
                throw new SecurityException("Unauthorized: Invalid token");
            }
            Object loginIdByToken = StpUtil.getLoginIdByToken(token);

            if(loginIdByToken == null) {
                throw new SecurityException("Unauthorized: Invalid token");
            }
            // 绑定用户ID到会话
            String userId = String.valueOf(loginIdByToken);

            accessor.setUser(() -> userId);
            SessionKit.addSession(sessionId, Long.valueOf(userId));

        }
        return message;
    }

    @Override
    public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
        ChannelInterceptor.super.postSend(message, channel, sent);
    }

    @Override
    public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
        ChannelInterceptor.super.afterSendCompletion(message, channel, sent, ex);
    }
}
